import time
import pandas as pd
from langchain_experimental.agents import create_pandas_dataframe_agent
from llms.llms_ev import ChatZhipuAI
from db.models.llm_profile.crud import get_llm_model_by_id

# 设置 API 密钥
api_key = get_llm_model_by_id(llm_model_id=1)["api_key"]
# 读取 CSV 文件
df = pd.read_csv('../data/total.csv')

print(df.head())

# 初始化模型
start_time = time.time()
try:
    llms = ChatZhipuAI(model="GLM-4-0520", api_key=api_key)
    print("模型初始化成功")
    elapsed_time_ms = (time.time() - start_time) * 1000  # 转换为毫秒
    print(f"模型初始化耗时：{elapsed_time_ms:.2f} ms")  # 输出保留两位小数
except Exception as e:
    print("模型初始化失败:", e)
    exit()

# 创建 Pandas DataFrame 代理
start_time = time.time()
try:
    agent_executor = create_pandas_dataframe_agent(
        llms,
        df,
        agent_type="openai-tools",
        verbose=True,
        allow_dangerous_code=True  # 允许危险代码执行
    )
    print("代理创建成功")
    elapsed_time_ms = (time.time() - start_time) * 1000  # 转换为毫秒
    print(f"代理创建耗时：{elapsed_time_ms:.2f} ms")  # 输出保留两位小数
except Exception as e:
    print("代理创建失败:", e)
    exit()

# 执行查询
start_time = time.time()
try:
    result = agent_executor.invoke({"input": "给出1982年同环比统计的数据的sql脚本"})
    print("查询成功")
    elapsed_time_ms = (time.time() - start_time) * 1000  # 转换为毫秒
    print(f"查询耗时：{elapsed_time_ms:.2f} ms")  # 输出保留两位小数
    print(result)
except Exception as e:
    print("查询失败:", e)
